package com.zdb.datastructures.tree.rbt;

public class RedBlackTree {

    private final static int RED = 1;
    private final static int BLACK = 0;

    private Node root;

    public static RedBlackTree genTestTree() {
        RedBlackTree redBlackTree = new RedBlackTree();

        redBlackTree.root = new Node(BLACK, 19, null);

        Node node5 = new Node(RED, 5, null);
        Node node7 = new Node(RED, 7, null);
        Node node13 = new Node(RED, 13, null);
        Node node35 = new Node(RED, 35, null);

        Node node1 = new Node(BLACK, 1, null);
        Node node12 = new Node(BLACK, 12, null);
        Node node30 = new Node(BLACK, 30, null);

        redBlackTree.root.setLeft(node5);
        redBlackTree.root.setRight(node30);

        node5.setLeft(node1);
        node5.setRight(node12);

        node12.setLeft(node7);
        node12.setRight(node13);

        node30.setRight(node35);


        return redBlackTree;
    }

    @Override public String toString() {
        return "RedBlackTree{" +
                "root=" + root +
                '}';
    }

    static class Node {
        int color = RED;
        int data;
        Node left;
        Node right;
        Node parent;

        public Node(int color, int data, Node parent) {
            this.color = color;
            this.data = data;
            this.parent = parent;
        }

        @Override public String toString() {

            StringBuilder sb = new StringBuilder("{"
                    + data + (color==RED?"R":"B"));

            if(left != null || right != null) {
                sb.append(":[");
                if(left != null) {
                    sb.append("L").append(left);
                }
                sb.append(",");
                if(right != null) {
                    sb.append("R").append(right);
                }

                sb.append("]");
            }


            sb.append("}");

            return sb.toString();
        }

        public boolean isRed() {
            return color == RED;
        }

        public boolean isBlack() {
            return color == BLACK;
        }

        public Node getUncle() {
            if(this.parent == null || this.parent.parent == null) {
                return null;
            }
            Node grandFather = this.parent.parent;
            if(grandFather.left == parent) {
                return grandFather.right;
            } else {
                return grandFather.left;
            }
        }

        public void setLeft(Node node) {
            this.left = node;
            if(node != null) {
                node.parent = this;
            }
        }

        public void setRight(Node node) {
            this.right = node;
            if(node != null) {
                node.parent = this;
            }
        }

        public boolean isParentLeftNode() {
            return parent.left == this;
        }

        public boolean isParentRightNode() {
            return parent.right == this;
        }

        /**
         * 以当前节点进行左旋，
         *
         * 1. 设置临时变量tmpNode保存当前节点的右子节点,parentNode保存当前节点的parent
         * 2. 设置当前节点的右子节点为当前节点的右子节点的左子节点
         * 3. 设置tmpNode的左子节点为当前节点
         * 4. 修改当前节点的parent对当前节点的指向，该为指向tmpNode
         *  4.1. 如果当前节点为根节点，那么将rbTree的root指向当前节点。
         *  4.2. 判断当前节点是parent的左子节点还是右子节点，
         *      对应修改parent的左子节点或者右子节点为tmpNode
         *
         *
         * @param rbtree 如果当前节点为根节点时需要要用到
         */
        public void rotateLeft(RedBlackTree rbtree) {
            Node parentNode = parent;

            // 1
            Node tmpNode = right;
            // 2
            this.setRight(right.left);;
            // 3
            tmpNode.setLeft(this);
            // 4
            setCurrentNodeParentNewChildNode(rbtree, parentNode, tmpNode);
        }

        /**
         * 4. 修改当前节点的parent对当前节点的指向，该为指向tmpNode
         *  4.1. 如果当前节点为根节点，那么将rbTree的root指向当前节点。
         *  4.2. 判断当前节点是parent的左子节点还是右子节点，
         *      对应修改parent的左子节点或者右子节点为tmpNode
         * @param rbtree
         * @param newChildNode
         */
        private void setCurrentNodeParentNewChildNode(
                RedBlackTree rbtree, Node parentNode, Node newChildNode) {
            if(parentNode == null) {
                rbtree.root = newChildNode;
                newChildNode.parent = null;
            } else {
                if(parentNode.left == this) {
                    parentNode.setLeft(newChildNode);
                } else {
                    parentNode.setRight(newChildNode);
                }
            }
        }

        /**
         * 以当前节点进行右旋
         *
         * 1. 设置临时变量tmpNode保存当前节点的左子节点,parentNode保存当前节点的parent
         * 2. 设置当前节点的左子节点为当前节点的左子节点的右子节点
         * 3. 设置tmpNode的右子节点为当前节点
         * 4. 修改当前节点的parent对当前节点的指向，该为指向tmpNode
         *  4.1. 如果当前节点为根节点，那么将rbTree的root指向当前节点。
         *  4.2. 判断当前节点是parent的左子节点还是右子节点，
         *      对应修改parent的左子节点或者右子节点为tmpNode
         *
         */
        public void rotateRight(RedBlackTree rbt) {

            // 1
            Node parentNode = parent;
            Node tmpNode = left;

            // 2
            setLeft(left.right);

            // 3
            tmpNode.setRight(this);

            // 4
            setCurrentNodeParentNewChildNode(rbt, parentNode, tmpNode);
        }
    }

    public void insert(int data) {
        if(root == null) {
            root = new Node(BLACK, data, null);
        } else {
            // 先按照bst插入节点
            Node insertNode = insert(root, data);
            // 修复违规情况
            fixRBTInsertion(insertNode, this);
        }
    }

    /**
     * 1. 先按照bst,插入节点
     * @param parent 不能为空
     * @param data
     */
    private Node insert(Node parent, int data) {
        // 小于的放到左边，大于等于的放到右边
        if(data < parent.data) {
            if(parent.left != null) {
                // 递归插入
                return insert(parent.left, data);
            } else {
                // 左边为空，直接设置左边节点
                parent.setLeft(new Node(RED, data, parent));
                return parent.left;
            }
        } else {
            if(parent.right != null) {
                // 递归插入
                return insert(parent.right, data);
            } else {
                parent.setRight(new Node(RED, data, parent));
                return parent.right;
            }
        }
    }

    /**
     * 在currentNode判断是否违规，如果违规，修复
     *
     * 正常情况：
     * case1.0 parent为黑色节点
     *
     * 违规情况：
     * case2.1. parent为红色节点，uncle为红色节点：
     *  > 1) 将parent,uncle都变成黑色
     *  > 2）将grandfather变成红色
     *  > 3) 将grandfather递归调用 fixRBTInsertion(grandfather)
     *
     * case2.2. parent为红色节点，uncle为黑色节点，且当前节点为parent的右子树：
     * > 左旋：
     * 1）以parent进行左旋
     *
     * case2.3. parent为红色节点，uncle为黑色节点，且当前节点为parent的左子树：
     * > 右旋
     * 1) 将parent设置为黑色，将grandfather设置为红色
     * 2) 以grandfather进行右旋
     *
     * @param currentNode 一定是红色节点
     */
    private void fixRBTInsertion(Node currentNode, RedBlackTree rbTree) {

        if(currentNode == null) {
        } else

        // 如果当前节点是根节点，直接返回
        if(currentNode.parent == null || rbTree.root == currentNode) {
        } else

        // case 1.0
        if(currentNode.parent.isBlack()) {
        } else {

            Node uncle = currentNode.getUncle();
            Node grandfather = currentNode.parent.parent;
            // case 2.1 uncle为空认为uncle是黑的
            if (currentNode.parent.isRed() && (uncle != null && uncle.isRed())) {
                setNodeBlack(currentNode.parent);
                setNodeBlack(uncle);
                setNodeRed(grandfather);
                fixRBTInsertion(grandfather, rbTree);
            } else
            // case 2.2
            if (currentNode.parent.isRed()
                    && (uncle == null || uncle.isBlack())
                    && currentNode.isParentRightNode()) {
                Node _parent = currentNode.parent;
                _parent.rotateLeft(rbTree);
                // 递归
                fixRBTInsertion(_parent, rbTree);
            } else
            // case 2.3
            if (currentNode.parent.isRed()
                    && (uncle == null || uncle.isBlack())
                    && currentNode.isParentLeftNode()) {

                if (grandfather == null) {
                    System.out.println("---");
                }

                // 设置parent为黑色，grandfather为红色
                setNodeBlack(currentNode.parent);
                setNodeRed(grandfather);

                // 以grandfather右旋
                grandfather.rotateRight(rbTree);
                // 递归
                fixRBTInsertion(grandfather, rbTree);

            }
        }

        rbTree.root.color = BLACK;
    }

    private void setNodeBlack(Node node) {
        if(node != null) {
            node.color = BLACK;
        }
    }

    private void setNodeRed(Node node) {
        if(node != null) {
            node.color = RED;
        }
    }
}
